
const WebSocket = require('ws');
const OpenAI = require('openai');
// 创建一个 WebSocket 服务器，监听在指定端口
const wss = new WebSocket.Server({ port: 39000 });

const apiKey = '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';

const openai = new OpenAI({ apiKey:apiKey });

const apiKey1 = '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';

const openai1 = new OpenAI({ apiKey:apiKey1 });

const modelList=[
    'ft:gpt-3.5-turbo-1106:personal::8LTW7fUm',
    'ft:gpt-3.5-turbo-1106:personal::8KHytL0m',
    'ft:gpt-3.5-turbo-1106:personal::8Nx5QLgv',
    'ft:gpt-3.5-turbo-1106:personal::8Q8orAtu',
    'ft:gpt-3.5-turbo-1106:personal::8SG7OH1b'
]

// 当有客户端连接时执行
wss.on('connection', (ws) => {
    console.log('Client connected');

    let stopFlag = false

    async function chatCompletions(recMsg) {

        ws.send(`openAIWSSBack`);
        let completion = {};
        if(modelList.includes(recMsg.model)){
            completion = await openai.chat.completions.create(recMsg);

        }else{
            completion = await openai1.chat.completions.create(recMsg);

        }
        console.log("completion:"+JSON.stringify(completion));
        ws.send(`openAIWSSStart`);
        
        for await (const chunk of completion) {
            console.log("chunk:"+JSON.stringify(chunk));
            console.log(chunk.choices[0].delta.content);

                    // 发送消息给客户端
            ws.send(chunk.choices[0].delta.content);

            if(stopFlag)break;

        }

        ws.send(`openAIWSSEnd`);
    }

    // 当收到消息时执行
    ws.on('message', (message) => {
        console.log(`Received: ${message}`);
        ws.send(`openAIWSSRecived`);
        if(message=='connectOpenAIWSS'){
            return
        }

        let recMsg= JSON.parse(message); 
        chatCompletions(recMsg);

    });

    // 当客户端断开连接时执行
    ws.on('close', (res) => {
        console.log('Client disconnected:'+JSON.stringify(res));
        if(res===3333){
            stopFlag = true
        }
    });
});